{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "from nbfinder import NotebookFinder\n",
    "sys.meta_path.append(NotebookFinder())\n",
    "import numpy as np\n",
    "import time\n",
    "from label_loader import  make_labels_for_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def make_yolo_masks_for_dataset( camfile_name, kwargs, labels_csv_file):\n",
    "        box_list = make_labels_for_dataset(camfile_name, labels_csv_file)\n",
    "        yolo_mask = create_yolo_gr_truth(box_list, kwargs)\n",
    "        return yolo_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_gr_truth_configs(kwargs):\n",
    "    scale_factor = float(kwargs[\"scale_factor\"])\n",
    "    xdim, ydim = kwargs[\"xdim\"], kwargs[\"ydim\"]\n",
    "    num_classes = kwargs[\"num_classes\"]\n",
    "    #make sure xy coords divide cleanly with scale_factor\n",
    "    assert xdim % scale_factor == 0 and ydim % scale_factor == 0, \"scale factor %i must divide the xy (%i, %i) coords cleanly \" %(scale_factor,xdim, ydim)\n",
    "    \n",
    "    xlen, ylen = xdim / int(scale_factor), ydim / int(scale_factor)\n",
    "    \n",
    "    #x,y,w,h,conf1,conf2 plus num_classes for one hot encoding\n",
    "    last_dim = 6 + num_classes\n",
    "    return scale_factor, xlen, ylen, last_dim, num_classes \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_xy_inds(x,y, scale_factor):\n",
    "        # get the indices to the lower left corner of the grid\n",
    "        \n",
    "        #scale x and y down\n",
    "        xs, ys = x / scale_factor, y / scale_factor\n",
    "        eps = 10*np.finfo(float).eps\n",
    "        #take the floor of x and y -> which is rounding to nearest bottom left corner\n",
    "        x_ind, y_ind = [np.floor(k - 10*eps ).astype('int') for k in [xs,ys]]\n",
    "        return x_ind, y_ind\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_xy_offsets(x,y, x_ind, y_ind, scale_factor):\n",
    "    #scale x and y down\n",
    "    xs, ys = x / scale_factor, y / scale_factor\n",
    "    \n",
    "    #get the offsets by subtracting the scaled lower left corner coords from the scaled real coords\n",
    "    xoff, yoff = xs - x_ind, ys - y_ind\n",
    "    \n",
    "    return xoff, yoff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_parametrized_wh(w,h,scale_factor):\n",
    "    ws , hs = w / scale_factor, h/ scale_factor\n",
    "    wp, hp = np.log2(ws), np.log2(hs)\n",
    "    return wp, hp\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def convert_class_to_one_hot(class_num, num_classes):\n",
    "    vec = num_classes * [0]\n",
    "    vec[class_num - 1] = 1\n",
    "    return vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_box_vector(coords, scale_factor, num_classes):\n",
    "    x,y,w,h,cls = coords\n",
    "    xind, yind = get_xy_inds(x,y,scale_factor)\n",
    "    xoff, yoff = get_xy_offsets(x, y, xind, yind, scale_factor)\n",
    "    wp, hp = get_parametrized_wh(w, h, scale_factor)\n",
    "    class_1hot_vec = convert_class_to_one_hot(cls, num_classes=num_classes)\n",
    "    objectness = [1, 0]\n",
    "    box_loc = [xoff, yoff, wp, hp]\n",
    "    box_vec = np.asarray(box_loc + objectness + class_1hot_vec)\n",
    "    return box_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_default_no_object_1hot(gr_truth):\n",
    "    #make the 5th number 1, so the objectness by defualt is 0,1 -> denoting no object\n",
    "    gr_truth[:,:,:,5] = 1.\n",
    "    return gr_truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def create_yolo_gr_truth(bbox_list, kwargs, caffe_format=False):\n",
    "        scale_factor, xlen, ylen, last_dim, num_classes = get_gr_truth_configs(kwargs)\n",
    "        \n",
    "        num_time_steps = len(bbox_list)\n",
    "        \n",
    "        gr_truth = np.zeros(( num_time_steps, xlen, ylen, last_dim ))\n",
    "        gr_truth = make_default_no_object_1hot(gr_truth)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        for time_step in range(num_time_steps):\n",
    "            for coords in bbox_list[time_step]:\n",
    "                x,y,w,h,cls = coords\n",
    "\n",
    "                xind, yind = get_xy_inds(x,y,scale_factor)\n",
    "                box_vec = get_box_vector(coords, scale_factor, num_classes)\n",
    "                gr_truth[time_step,xind,yind,:] = box_vec\n",
    "\n",
    "        if caffe_format:\n",
    "            gr_truth = np.transpose(gr_truth, axes=(0,3,1,2))\n",
    "        return gr_truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def test_grid(bbox, grid, xdim, ydim, scale_factor,num_classes, caffe_format=False):\n",
    "    scale_factor = float(scale_factor)\n",
    "    cls = int(bbox[4])\n",
    "    x,y = bbox[0] / scale_factor, bbox[1] / scale_factor\n",
    "\n",
    "    xo,yo = x - np.floor(x), y - np.floor(y)\n",
    "    w,h = np.log2(bbox[2] / scale_factor), np.log2(bbox[3] / scale_factor)\n",
    "\n",
    "\n",
    "    depth = 6 + num_classes\n",
    "    if caffe_format:\n",
    "        l_box = grid[:depth,x,y]\n",
    "    else:\n",
    "        l_box = grid[int(x),int(y),:depth]\n",
    "\n",
    "    lbl = num_classes*[0]\n",
    "    lbl[cls-1] = 1\n",
    "    \n",
    "    real_box = [xo,yo,w,h,1.,0.]\n",
    "    real_box.extend(lbl)\n",
    "    \n",
    "    print l_box\n",
    "    print real_box\n",
    "    assert np.allclose(l_box, real_box), \"Tests Failed\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.75      0.796875 -1.830075 -1.830075  1.        0.        0.        1.\n",
      "  0.        0.      ]\n",
      "[0.75, 0.796875, -1.8300749985576876, -1.8300749985576876, 1.0, 0.0, 0, 1, 0, 0]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python2.7/dist-packages/ipykernel/__main__.py:12: VisibleDeprecationWarning: using a non-integer number instead of an integer will result in an error in the future\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    kwargs = {  \"metadata_dir\": \"/home/evan/data/climate/labels/\",\n",
    "                \"scale_factor\": 64, \n",
    "                \"xdim\":768,\n",
    "                \"ydim\":1152,\n",
    "                \"time_steps_per_file\": 8,\n",
    "                \"num_classes\": 4, }\n",
    "\n",
    "    camfile_name = \"/home/evan/data/climate/input/cam5_1_amip_run2.cam2.h2.1979-01-05-00000.nc\"\n",
    "    labels_csv_file = \"/home/evan/data/climate/labels/labels.csv\"\n",
    "    ym = make_yolo_masks_for_dataset(camfile_name,\n",
    "                                     kwargs,\n",
    "                                    labels_csv_file)\n",
    "    box_list = make_labels_for_dataset(camfile_name, labels_csv_file)\n",
    "    box = box_list[0][0]\n",
    "    box_vec = get_box_vector(box, scale_factor=64., num_classes=4)\n",
    "    test_grid(box, ym[0], kwargs[\"xdim\"], kwargs[\"ydim\"], kwargs[\"scale_factor\"], kwargs[\"num_classes\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
